import torch
from torch import nn

IMG_SIZE = 32
LR = 0.9

class LeNet(nn.Module):
    """LeNet-5"""
    def __init__(self, num_classes, in_channels=3):
        super(LeNet, self).__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, 6, kernel_size=5, stride=1, padding=0), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0), nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 120, kernel_size=5, stride=1, padding=0), nn.ReLU(),
            nn.Flatten(),
            nn.Linear(120, 84), nn.ReLU(),
            nn.Linear(84, num_classes)
        )
    

    def forward(self, input):
        return self.net(input)


    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR